#!/usr/bin/env python
# encoding: utf-8
# @author: lipeng
# functionality: parse xml to txt that contains object class and bboxes
# format: imagepath x1 y1 x2 y2 cls_idx ...
import os
from tqdm import tqdm
import xml.etree.ElementTree as ET

VOC_CLASSES = (    # always index 0
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')

def parse_xml(xml_file):
    tree = ET.parse(xml_file)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        difficulty = int(obj.find('difficult').text)
        if 1 == difficulty:
            continue
        obj_struct['name'] = obj.find('name').text
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
                              int(float(bbox.find('ymin').text)),
                              int(float(bbox.find('xmax').text)),
                              int(float(bbox.find('ymax').text))]
        objects.append(obj_struct)
    return objects

def get_object_bboxes(imgnames_file, root_path, saved_file):

    txt_file = open(saved_file,'w')

    test_file = open(imgnames_file,'r')
    lines = test_file.readlines()
    lines = [x.strip().split('/')[-1][:-4] for x in lines]

    VOC_Annotations = os.path.join(root_path, 'Annotations')
    VOC_JPEGImages = os.path.join(root_path, 'JPEGImages')
    xml_files = os.listdir(VOC_Annotations)

    for xml_file in tqdm(xml_files):
        if xml_file.split('.')[0] not in lines:
            continue
        image_path = os.path.join(VOC_JPEGImages, xml_file.split('.')[0] + '.jpg')
        results = parse_xml(os.path.join(VOC_Annotations, xml_file))
        if len(results)==0:
            print(xml_file)
            continue
        txt_file.write(image_path)
        for result in results:
            class_name = result['name']
            bbox = result['bbox']
            class_idx = VOC_CLASSES.index(class_name)
            txt_file.write(' '+str(bbox[0])+' '+str(bbox[1])+' '+str(bbox[2])+' '+str(bbox[3])+' '+str(class_idx))
        txt_file.write('\n')
    txt_file.close()

if __name__=="__main__":
    imgnames_file = 'voc2007test.txt'
    root_path = '/media/tclwh/lipeng/dataset/VOCdevkit/VOC2007/'
    saved_file = 'voc2007imgtest.txt'
    get_object_bboxes(imgnames_file, root_path, saved_file)
    print('done!')
